#!/usr/bin/env python3
"""
演示脚本：展示 MoonBit 重构后的人脸关键点检测功能
这个脚本提供了一个 Python 版本的演示，展示重构后的功能
"""

import cv2
import dlib
import numpy as np
import os
import sys

class FaceLandmarksDemo:
    def __init__(self, model_path="shape_predictor_81_face_landmarks.dat"):
        """初始化人脸关键点检测演示"""
        self.model_path = model_path
        self.detector = None
        self.predictor = None
        self.cap = None
        self.writer = None
        
    def init_detector(self):
        """初始化 dlib 人脸检测器和关键点预测器"""
        try:
            print("正在初始化人脸检测器...")
            self.detector = dlib.get_frontal_face_detector()
            
            if not os.path.exists(self.model_path):
                print(f"错误: 模型文件 {self.model_path} 不存在")
                return False
                
            self.predictor = dlib.shape_predictor(self.model_path)
            print("人脸检测器初始化完成")
            return True
        except Exception as e:
            print(f"初始化人脸检测器失败: {e}")
            return False
    
    def init_camera(self, device_id=0, width=1280, height=720):
        """初始化摄像头"""
        try:
            print("正在初始化摄像头...")
            self.cap = cv2.VideoCapture(device_id)
            
            if not self.cap.isOpened():
                print("错误: 无法打开摄像头")
                return False
            
            # 设置分辨率
            self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, width)
            self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, height)
            
            print(f"摄像头初始化完成 - 分辨率: {width}x{height}")
            return True
        except Exception as e:
            print(f"初始化摄像头失败: {e}")
            return False
    
    def init_video_writer(self, output_file="output.avi", fps=20.0):
        """初始化视频写入器"""
        try:
            # 获取实际的摄像头分辨率
            width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
            height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
            
            fourcc = cv2.VideoWriter_fourcc(*'XVID')
            self.writer = cv2.VideoWriter(output_file, fourcc, fps, (width, height))
            
            print(f"视频写入器初始化完成 - 输出文件: {output_file}")
            return True
        except Exception as e:
            print(f"初始化视频写入器失败: {e}")
            return False
    
    def detect_and_draw_landmarks(self, frame):
        """检测人脸并绘制关键点"""
        try:
            # 转换为灰度图
            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            
            # 检测人脸
            faces = self.detector(gray)
            
            # 对每个检测到的人脸绘制关键点
            for face in faces:
                # 预测关键点
                landmarks = self.predictor(gray, face)
                
                # 绘制81个关键点
                for i in range(81):
                    point = landmarks.part(i)
                    cv2.circle(frame, (point.x, point.y), 2, (0, 255, 0), -1)
            
            return frame, len(faces)
        except Exception as e:
            print(f"人脸检测错误: {e}")
            return frame, 0
    
    def run(self):
        """运行主循环"""
        if not self.init_detector():
            return False
        
        if not self.init_camera():
            return False
        
        if not self.init_video_writer():
            return False
        
        print("\n=== MoonBit 人脸关键点检测演示 ===")
        print("功能展示:")
        print("- 实时人脸检测")
        print("- 81点关键点预测")
        print("- 视频录制")
        print("- 按 'q' 键退出")
        print("=====================================\n")
        
        frame_count = 0
        face_count_total = 0
        
        try:
            while True:
                # 读取摄像头帧
                ret, frame = self.cap.read()
                if not ret:
                    print("无法读取摄像头帧")
                    break
                
                # 水平翻转（镜像效果）
                frame = cv2.flip(frame, 1)
                
                # 检测人脸并绘制关键点
                processed_frame, face_count = self.detect_and_draw_landmarks(frame)
                
                # 显示处理后的帧
                cv2.imshow('MoonBit Face Landmarks Detection Demo', processed_frame)
                
                # 写入视频文件
                if self.writer:
                    self.writer.write(processed_frame)
                
                # 统计信息
                frame_count += 1
                face_count_total += face_count
                
                if frame_count % 30 == 0:  # 每30帧显示一次统计
                    avg_faces = face_count_total / frame_count
                    print(f"已处理 {frame_count} 帧, 平均检测到 {avg_faces:.1f} 个人脸")
                
                # 检查退出键
                key = cv2.waitKey(1) & 0xFF
                if key == ord('q'):
                    break
                    
        except KeyboardInterrupt:
            print("\n用户中断程序")
        except Exception as e:
            print(f"运行时错误: {e}")
        finally:
            self.cleanup()
        
        print(f"\n程序结束 - 总共处理了 {frame_count} 帧")
        return True
    
    def cleanup(self):
        """清理资源"""
        print("正在清理资源...")
        
        if self.cap:
            self.cap.release()
        
        if self.writer:
            self.writer.release()
        
        cv2.destroyAllWindows()
        print("资源清理完成")

def main():
    """主函数"""
    print("MoonBit 人脸关键点检测项目演示")
    print("这个演示展示了重构后的功能特性\n")
    
    # 检查模型文件
    model_path = "shape_predictor_81_face_landmarks.dat"
    if not os.path.exists(model_path):
        print(f"错误: 找不到模型文件 {model_path}")
        print("请确保模型文件在当前目录中")
        return 1
    
    # 创建并运行演示
    demo = FaceLandmarksDemo(model_path)
    
    if demo.run():
        print("演示完成!")
        return 0
    else:
        print("演示失败!")
        return 1

if __name__ == "__main__":
    sys.exit(main())